#!/bin/bash
# EbbFlow 配置文件和密钥生成脚本（支持从环境变量读取IP）
# 使用方法: 
#   方式1（使用环境变量）: export IP_LIST="ip1,ip2,ip3,ip4" && ./run.sh
#   方式2（直接运行）: ./run.sh

set -e

# ==================== 颜色定义 ====================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# ==================== 日志函数 ====================
log_step() {
    echo -e "${BLUE}==>${NC} ${GREEN}$1${NC}"
}

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

# ==================== 配置参数 ====================
# 项目配置
PROJECT_ROOT="/home/ebbflow"
EC2_DIR="${PROJECT_ROOT}/ec2"
BIN_DIR="${PROJECT_ROOT}/bin"
ETC_DIR="${PROJECT_ROOT}/etc"

# AWS 凭证
AWS_ACCESS_KEY="${AWS_ACCESS_KEY_ID:-}"
AWS_SECRET_KEY="${AWS_SECRET_ACCESS_KEY:-}"

# 从环境变量读取 IP 列表（如果存在）
if [ -n "$IP_LIST" ]; then
    log_info "检测到环境变量 IP_LIST: $IP_LIST"
    # 将逗号分隔的IP转换为数组
    IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
    N_NODES=${#IP_ARRAY[@]}
    log_info "自动检测节点数: $N_NODES"
else
    log_warning "未检测到 IP_LIST 环境变量，使用默认配置"
    N_NODES=4
    IP_ARRAY=()
fi

# 实验配置
CONSENSUS_TYPE=14
# BATCH_SIZE=${BATCH_SIZE:-30000}  # 支持从环境变量读取
BATCH_SIZE=5000
TX_SIZE=250
# PORT_BASE=${PORT_BASE:-8000}      # 支持从环境变量读取
PORT_BASE=8000

# K_THRESHOLD 自动计算：对于 n 个节点，k = floor((n+1)/2)
# 这样可以容忍 f = floor((n-1)/3) 个拜占庭节点
if [ $N_NODES -eq 4 ]; then
    K_THRESHOLD=2  # 对于4个节点，推荐k=2
elif [ $N_NODES -ge 3 ]; then
    K_THRESHOLD=$(( (N_NODES + 1) / 2 ))
else
    K_THRESHOLD=1
fi

TOTAL_NODES=$N_NODES

# ==================== 检查环境 ====================
check_environment() {
    log_step "检查运行环境..."
    
    # 检查 Python3
    if ! command -v python3 &> /dev/null; then
        log_error "未找到 python3，请先安装 Python 3"
        exit 1
    fi
    
    # 检查 keygen 工具
    if [ ! -f "$BIN_DIR/keygen" ]; then
        log_error "未找到 $BIN_DIR/keygen 工具"
        exit 1
    fi
    
    # 检查 keygen 可执行权限
    if [ ! -x "$BIN_DIR/keygen" ]; then
        log_warning "keygen 没有执行权限，正在添加..."
        chmod +x "$BIN_DIR/keygen"
    fi
    
    # 检查项目目录
    if [ ! -d "$PROJECT_ROOT" ]; then
        log_error "项目目录 $PROJECT_ROOT 不存在"
        exit 1
    fi
    
    log_info "环境检查完成"
}

# ==================== 生成密钥 ====================
generate_keys() {
    log_step "生成协议密钥..."
    
    cd "$PROJECT_ROOT"
    
    # 创建密钥目录
    mkdir -p etc/thresprf_key
    
    # 生成门限 PRF 密钥
    log_info "生成门限 PRF 密钥 (n=$N_NODES, k=$K_THRESHOLD)..."
    ./bin/keygen $N_NODES $K_THRESHOLD
    
    # 检查密钥生成结果
    if [ ! -d "etc/thresprf_key" ] || [ -z "$(ls -A etc/thresprf_key)" ]; then
        log_error "密钥生成失败"
        exit 1
    fi
    
    # 显示生成的密钥文件
    log_info "生成的密钥文件："
    ls -lh etc/thresprf_key/
    
    log_info "密钥生成完成"
}

# ==================== 生成配置文件 ====================
generate_config() {
    log_step "生成配置文件..."
    
    cd "$PROJECT_ROOT"
    mkdir -p etc
    
    # 构建 replicas 数组
    local replicas_json="[]"
    if [ ${#IP_ARRAY[@]} -gt 0 ]; then
        replicas_json="["
        for i in "${!IP_ARRAY[@]}"; do
            local ip="${IP_ARRAY[$i]}"
            local port=$((PORT_BASE + i))
            if [ $i -gt 0 ]; then
                replicas_json+=","
            fi
            replicas_json+=$(cat <<EOF

        {
            "id": $i,
            "ip": "$ip",
            "port": $port
        }
EOF
)
        done
        replicas_json+="
    ]"
    fi
    
    # 创建 conf.json
    cat > etc/conf.json << EOF
{
    "consensus": $CONSENSUS_TYPE,
    "batch_size": $BATCH_SIZE,
    "tx_size": $TX_SIZE,
    "n_nodes": $N_NODES,
    "k_threshold": $K_THRESHOLD,
    "replicas": $replicas_json
}
EOF
    
    log_info "配置文件生成完成: etc/conf.json"
    echo ""
    cat etc/conf.json
    echo ""
}

# ==================== 显示配置摘要 ====================
show_summary() {
    echo ""
    echo -e "${BLUE}==================== 配置摘要 ====================${NC}"
    echo -e "${GREEN}节点数量:${NC} $N_NODES"
    echo -e "${GREEN}门限参数:${NC} $K_THRESHOLD"
    echo -e "${GREEN}共识类型:${NC} $CONSENSUS_TYPE"
    echo -e "${GREEN}批处理大小:${NC} $BATCH_SIZE"
    echo -e "${GREEN}交易大小:${NC} $TX_SIZE"
    echo -e "${GREEN}起始端口:${NC} $PORT_BASE"
    
    if [ ${#IP_ARRAY[@]} -gt 0 ]; then
        echo -e "${GREEN}节点列表:${NC}"
        for i in "${!IP_ARRAY[@]}"; do
            echo "  节点 $i: ${IP_ARRAY[$i]}:$((PORT_BASE + i))"
        done
    fi
    echo -e "${BLUE}=================================================${NC}"
    echo ""
}

# ==================== 主函数 ====================
main() {
    echo -e "${GREEN}=================================${NC}"
    echo -e "${GREEN}EbbFlow 配置文件和密钥生成工具${NC}"
    echo -e "${GREEN}=================================${NC}"
    echo ""
    
    # 检查环境
    check_environment
    
    # 显示配置
    show_summary
    
    # 生成密钥
    generate_keys
    
    # 生成配置文件
    generate_config
    
    echo ""
    log_step "全部完成！"
    echo -e "${GREEN}密钥目录:${NC} $PROJECT_ROOT/etc/thresprf_key"
    echo -e "${GREEN}配置文件:${NC} $PROJECT_ROOT/etc/conf.json"
    echo ""
    
    # 如果有IP列表，提示如何部署
    if [ ${#IP_ARRAY[@]} -gt 0 ]; then
        echo -e "${YELLOW}下一步操作建议:${NC}"
        echo "1. 将配置文件分发到各个节点"
        echo "2. 将密钥文件分发到各个节点（每个节点需要对应的密钥）"
        echo "3. 在各个节点上启动服务"
        echo ""
        echo -e "${YELLOW}分发命令示例:${NC}"
        if [ -n "$HOST_LIST" ]; then
            echo "  # 使用 HOST_LIST 环境变量"
            echo "  IFS=',' read -ra HOSTS <<< \"\$HOST_LIST\""
            echo "  for host in \"\${HOSTS[@]}\"; do"
            echo "    scp -r etc/conf.json \$host:/home/ubuntu/ebbflow/etc/"
            echo "    scp -r etc/thresprf_key \$host:/home/ubuntu/ebbflow/etc/"
            echo "  done"
        else
            echo "  for i in {0..$((N_NODES-1))}; do"
            echo "    scp -r etc/conf.json ubuntu@\${IP_ARRAY[\$i]}:/path/to/ebbflow/etc/"
            echo "    scp -r etc/thresprf_key ubuntu@\${IP_ARRAY[\$i]}:/path/to/ebbflow/etc/"
            echo "  done"
        fi
    fi
}

# 执行主函数
main